import torch
from matplotlib import pyplot as plt
import numpy as np


fig = plt.figure()
x = np.arange(-4, 4, 0.025)
plt.plot(x,x**2)
plt.title("y = x^2")
def f(x):
    return x**2
def h(x):
    return 2*x
# 初值
x = 4
# 动量与指数衰减的初值
m = 0
v = 0
# 动量与指数衰减的参数
β1 = 0.9
β2 = 0.999
# 学习率
η = 0.01
ε = 10e-8
iters = 0
X = []
Y = []
while iters<8000:
    iters+=1
    X.append(x)
    Y.append(f(x))
    m = β1*m + (1-β1)*h(x)
    v = β2*v + (1-β2)*h(x)**2
    m_het = m/(1-β1**iters)
    v_het = v/(1-β2**iters)
    x = x - η/np.sqrt(v_het+ε)*m_het
    print(iters,x)
plt.plot(X,Y,"ro")
plt.show()


params = list(m.parameters())
torch.optim.Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=0, amsgrad=False)